#include <src/pcps.h>
#include <src/input.h>
#include <src/hamiltonian.h>
#include <src/wavefunction.h>
#include <src/time_evolution.h>
#include <src/stochastic_optimization.h>
#include <src/projected_optimization_ti.h>
#include <src/fullci.h>
#include <src/pmc.h>
#include <src/mc_energy.h>

//--------------------------------------------------------------------------------
// pcps::driver -- controls which calculations to perform
//--------------------------------------------------------------------------------

template <class SCALAR>
void pcps::driver(const pcps::Input & userinp) {

  // get MPI info
  const MPI::Comm & comm = MPI::COMM_WORLD;
  const int nproc = comm.Get_size();
  const int myrank = comm.Get_rank();

  // create the hamiltonian
  const pcps::Hamiltonian<SCALAR> ham(userinp);

  // create the initial wavefunction
  pcps::Wavefunction<SCALAR> wfn(userinp, ham);

  // if requested, print the initial wavefunction
  if (myrank == 0 && userinp.verbose_print()) wfn.print(userinp);

  // projected cps optimization (via Newton Raphson and/or a local hamiltonian sweep)
  if (userinp.optimization() == pcps::keyword::Opt_Projected)
    pcps::projected_driver(userinp, ham, wfn);

  // translationally invariant projected cps optimization
  else if (userinp.optimization() == pcps::keyword::Opt_Projected_TI)
    pcps::projected_driver_ti(userinp, ham, wfn);

  // variational energy evaluation (no actual optimization)
  else if (userinp.optimization() == pcps::keyword::Opt_Energy_Evaluation)
    pcps::stochastic_optimization(userinp, ham, wfn);
    //const double e = pcps::real( pcps::mc_energy(userinp, ham, wfn) );
    //if (myrank == 0)
    //  std::cout << std::endl
    //            << boost::format("energy = %20.12f") % e
    //            << std::endl;

  // variational steepest descent optimization (via MC sampling)
  else if (userinp.optimization() == pcps::keyword::Opt_Steepest_Descent)
    pcps::stochastic_optimization(userinp, ham, wfn);

  // stochastic imaginary time evolution (via MC sampling)
  else if (userinp.optimization() == pcps::keyword::Opt_Stochastic_Imaginary_Time)
    pcps::stochastic_optimization(userinp, ham, wfn);

  // stochastic real time evolution (via MC sampling)
  else if (userinp.optimization() == pcps::keyword::Opt_Stochastic_Real_Time)
    pcps::stochastic_optimization(userinp, ham, wfn);

  // optimization by imaginary time evolution (via local cluster summations)
  else if (userinp.optimization() == pcps::keyword::Opt_Imaginary_Time)
    pcps::time_evolution_driver(userinp, ham, wfn, false);

  // real time evolution (via local cluster summations)
  else if (userinp.optimization() == pcps::keyword::Opt_Real_Time)
    pcps::time_evolution_driver(userinp, ham, wfn, true);

  // full configuration interaction
  else if (userinp.optimization() == pcps::keyword::Opt_Full_CI)
    pcps::fci::fullci<SCALAR>(userinp, ham);

  // projector monte carlo
  else if (userinp.optimization() == pcps::keyword::Opt_PMC)
    pcps::pmc::driver<SCALAR>(userinp, ham, wfn);

  // dmrg integral conversion
  else if (userinp.optimization() == pcps::keyword::Opt_Dump_DMRG_Ints)
    ham.write_dmrg_ints(userinp);

  // no optimization
  else if (userinp.optimization() == pcps::keyword::Opt_None)
  {}

  // unknown optimization
  else
    throw pcps::Exception( (boost::format("\"%s\" is not a supported optimization") % userinp.optimization()).str() );

  // convert correlators from spins to fermion and dump the wavefunction
  if (userinp.spin_to_fermion_dump()) {

    // create a vector to hold converted correlators
    std::vector< pcps::Correlator<SCALAR> > f_corr_vec;

    // convert each correlator
    for (typename std::vector< pcps::Correlator<SCALAR> >::iterator corr = wfn.correlators().begin(); corr != wfn.correlators().end(); corr++) {

      // get the fermion lattice sites for the converted correlator
      std::vector<int> fsites;
      for (std::vector<int>::const_iterator site = corr->site_vec().begin(); site != corr->site_vec().end(); site++) {
        fsites.push_back(2*(*site));
        fsites.push_back(2*(*site)+1);
      }

      // create the fermion correlator
      pcps::Correlator<SCALAR> fcorr(fsites, 2, pcps::unity<SCALAR>(), corr->fixed());

      // transfer the values to the new correlator
      pcps::Cluster fcluster(2*userinp.nsites(), 2, fsites);
      for (pcps::Cluster scluster(userinp.nsites(), userinp.configs_per_site(), corr->site_vec()); !scluster.finished(); scluster++) {
        for (std::vector<int>::const_iterator site = corr->site_vec().begin(); site != corr->site_vec().end(); site++) {
          fcluster.set_occ(2*(*site)+0, 1 - scluster.ket_occ(*site));
          fcluster.set_occ(2*(*site)+1, scluster.ket_occ(*site));
        }
        fcorr[fcluster] = (*corr)[scluster];
      }

      // add the fermionic correlator
      f_corr_vec.push_back(fcorr);

    }

    // exchange the spin and fermionic correlators
    std::swap(wfn.correlators(), f_corr_vec);

    // write the fermionic wavefunction to disk
    wfn.dump_binary(userinp, "pcps_spin_to_fermion_wfn.bin");

    // put the spin correlators back in the wavefunction
    std::swap(wfn.correlators(), f_corr_vec);

  }

}

// explicitly instantiate the template
template void pcps::driver< double               >(const pcps::Input &);
template void pcps::driver< std::complex<double> >(const pcps::Input &);
